import numpy as np
import matplotlib.pyplot as plt
import matplotlib

matplotlib.rcParams['font.sans-serif'] = ['SimHei'] 
matplotlib.rcParams['axes.unicode_minus'] = False  

def assign_cluster(x, c):
    distances = np.linalg.norm(c - x, axis=1)
    y = np.argmin(distances)
    return y

def Kmeans(data, k, epsilon=1e-3, iteration=100):
    n_samples, n_features = data.shape
    np.random.seed(42)
    centroids = data[np.random.choice(n_samples, k, replace=False)]
    Y = np.zeros(n_samples, dtype=int)
    for _ in range(iteration):
        for i in range(n_samples):
            Y[i] = assign_cluster(data[i], centroids)
        new_centroids = np.array([data[Y == idx].mean(axis=0) for idx in range(k)])
        if np.linalg.norm(new_centroids - centroids) < epsilon:
            break
        centroids = new_centroids
    return Y, centroids

if __name__ == "__main__":
    np.random.seed(42)
    cluster1 = np.random.randn(30, 2) + [5, 5]
    cluster2 = np.random.randn(30, 2) + [0, 0]
    cluster3 = np.random.randn(40, 2) + [10, 0]
    test_data = np.vstack([cluster1, cluster2, cluster3])
    k = 3
    labels, centers = Kmeans(test_data, k=k)

    print(f"聚类完成！共{len(test_data)}个样本，聚为{k}类")
    print("前10个样本的簇标签:", labels[:10])
    print("最终聚类中心：\n", centers)

    plt.figure(figsize=(8, 6))
    for i in range(k):
        plt.scatter(test_data[labels == i, 0], test_data[labels == i, 1], label=f"簇{i+1}")
    plt.scatter(centers[:, 0], centers[:, 1], c="black", marker="*", s=200, label="聚类中心")
    plt.legend()
    plt.title("K-means聚类结果") 
    plt.show()
